#include <cmath>
#include <cassert>
#include <utility>
#include <iostream>

template <class K, class V>
struct AVLTreeNode
{
    // 需要parent指针，后续更新平衡因⼦可以看到
    std::pair<K, V> _kv;
    AVLTreeNode<K, V> *_left;
    AVLTreeNode<K, V> *_right;
    AVLTreeNode<K, V> *_parent;
    // balance factor 右子树的高度减去左子树的高度([-1,1])
    // 这个值不变说明子树的高度不变
    int _bf;
    AVLTreeNode(const std::pair<K, V> &kv)
        : _kv(kv),
          _left(nullptr),
          _right(nullptr),
          _parent(nullptr),
          _bf(0)
    {
    }
};
template <class K, class V>
class AVLTree
{
    using Node = AVLTreeNode<K, V>;

public:
    Node *getRoot() { return _root; }
    bool insert(const std::pair<K, V> &kv)
    {
        if (_root == nullptr)
        {
            _root = new Node(kv);
            return true;
        }
        // 找到该节点要插入的地方
        Node *parent = nullptr;
        Node *cur = _root;
        while (cur)
        {
            if (kv.first > cur->_kv.first)
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (kv.first < cur->_kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            // 如果k值相等，说明树中已经存在了，不能插入
            else
                return false;
        }

        // 找到了一个合适的插入位置
        cur = new Node(kv);

        // 决定插在左边还是右边
        if (kv.first > parent->_kv.first)
            parent->_right = cur;
        else
            parent->_left = cur;

        cur->_parent = parent;

        // 更新平衡因子

        while (parent)
        {
            if (parent->_left == cur)
                parent->_bf--;
            else
                parent->_bf++;

            if (parent->_bf == 0)
                break;

            else if (abs(parent->_bf) == 1)
            {
                cur = parent;
                parent = parent->_parent;
            }

            else if (abs(parent->_bf) == 2)
            {
                // 旋转
                // 左右旋写完更新一下这里
                if (parent->_bf == -2 && cur->_bf == -1)
                {
                    RotateRight(parent);
                }
                else if (parent->_bf == -2 && cur->_bf == 1)
                {
                    RotateLeftRight(parent);
                }
                else if (parent->_bf == 2 && cur->_bf == -1)
                {
                    RotateRightLeft(parent);
                }
                else if (parent->_bf == 2 && cur->_bf == 1)
                {
                    RotateLeft(parent);
                }
                else
                {
                    assert(false);
                }
                break;
            }

            else
            {
                assert(false);
            }
        }
        return true;
    }

    /**
     * 保持搜索树的规则
     * 变平衡
     * 降低树的高度
     */
    // 右单旋: 某个节点的左子树失衡，需要把他的左孩子提上来，解决左左失衡的问题
    void RotateRight(Node *parent)
    {
        /**
         * 假设某一个子树新插入节点后不平衡了，把这个子树的父亲节点旋转到根节点上
         * 细节：子节点的判空，非空才能使用，parent指针的控制
         * 判断这是单独的子树还是根开始的子树，根据parent指针判断
         * 提前记录parent指针
         * 平衡因子的更新，只有子树高度发生变化了才需要更新
         */

        Node *subL = parent->_left;
        Node *subLR = subL->_right;

        Node *parent_ = parent->_parent;

        parent->_left = subLR;
        if (subLR)
            subLR->_parent = parent;

        subL->_right = parent;
        parent->_parent = subL;

        if (parent_ == nullptr)
        {
            _root = subL;
            subL->_parent = nullptr;
        }
        else
        {
            // 判断这是左子树还是右子树
            if (parent_->_left == parent)
            {
                parent_->_left = subL;
            }
            else
            {
                parent_->_right = subL;
            }
            subL->_parent = parent_;
        }

        parent->_bf = 0;
        subL->_bf = 0;
    }

    void RotateLeft(Node *parent)
    {
        /**
         * 和右单旋逻辑一样
         */

        Node *subR = parent->_right;
        Node *subRL = subR->_left;

        Node *parent_ = parent->_parent;

        parent->_right = subRL;

        if (subRL)
            subRL->_parent = parent;

        subR->_left = parent;
        parent->_parent = subR;

        if (parent_ == nullptr)
        {
            _root = subR;
            subR->_parent = nullptr;
        }
        else
        {
            if (parent == parent_->_left)
            {
                parent_->_left = subR;
            }
            else
            {
                parent_->_right = subR;
            }
            subR->_parent = parent_;
        }
        parent->_bf = 0;
        subR->_bf = 0;
    }

    // 左右双旋
    void RotateLeftRight(Node *parent)
    {
        // 提前记录平衡因子
        Node *subL = parent->_left;
        Node *subLR = subL->_right;
        int bf = subLR->_bf;
        // parent->left左旋
        // parent右旋
        RotateLeft(parent->_left);
        RotateRight(parent);
        // 更新平衡因子
        // 主要的调整还是左旋和右旋的点
        if (bf == -1)
        {
            parent->_bf = 1;
            subL->_bf = 0;
            subLR->_bf = 0;
        }
        else if (bf == 1)
        {
            parent->_bf = 0;
            subL->_bf = -1;
            subLR->_bf = 0;
        }
        else if (bf == 0)
        {
            parent->_bf = 0;
            subL->_bf = 0;
            subLR->_bf = 0;
        }
        else
        {
            // 处理其他错误
            assert(false);
        }
    }

    // 右左双旋
    void RotateRightLeft(Node *parent)
    {
        // 提前记录平衡因子
        Node *subR = parent->_right;
        Node *subRL = subR->_left;
        int bf = subRL->_bf;
        // parent->right右旋
        // parent左旋
        // 注意旋转顺序
        RotateRight(parent->_right);
        RotateLeft(parent);
        // 更新平衡因子
        // 主要的调整还是左旋和右旋的点
        if (bf == -1)
        {
            parent->_bf = 0;
            subR->_bf = 1;
            subRL->_bf = 0;
        }
        else if (bf == 1)
        {
            parent->_bf = -1;
            subR->_bf = 0;
            subRL->_bf = 0;
        }
        else if (bf == 0)
        {
            parent->_bf = 0;
            subR->_bf = 0;
            subRL->_bf = 0;
        }
        else
        {
            // 处理其他错误
            assert(false);
        }
    }

    Node *find(const K &key)
    {
        Node *cur = _root;
        while (cur)
        {
            if (cur->_kv.first > key)
            {
                cur = cur->_left;
            }
            else if (cur->_kv.first < key)
            {
                cur = cur->_right;
            }
            else
                return cur;
        }
        return nullptr;
    }

    bool IsBalace(Node *root)
    {
        if (!root)
            return true;

        int leftHigh = getTreeHigh(root->_left);
        int rightHigh = getTreeHigh(root->_right);

        int diff = rightHigh - leftHigh;

        if (abs(rightHigh - leftHigh) >= 2)
        {
            std::cout << root->_kv.first << "balabce exception" << std::endl;
            return false;
        }
        if (root->_bf != diff)
        {
            std::cout << root->_kv.first << "balabce exception" << std::endl;
            return false;
        }
        return IsBalace(root->_left) && IsBalace(root->_right);
    }

    void InOrder(Node *root)
    {
        if (!root)
            return;
        InOrder(root->_left);
        std::cout << root->_kv.first << ":" << root->_kv.second << std::endl;
        InOrder(root->_right);
    }

private:
    int getTreeHigh(Node *root)
    {
        if (!root)
            return 0;
        int leftHigh = getTreeHigh(root->_left);
        int rightHigh = getTreeHigh(root->_right);

        return leftHigh > rightHigh ? (leftHigh + 1) : (rightHigh + 1);
    }

private:
    Node *_root = nullptr;
};
